import gradio as gr
from PIL import Image, ImageOps, ImageChops
from utils import generate_random_dir, generate_random_file
from utils.ps_canvas import PSTools
from utils.maas import MaaS
from portrait import PortraitSolution
import os
import requests
from loguru import logger
from io import BytesIO
from pydantic import BaseModel
from requests_toolbelt.multipart.encoder import MultipartEncoder
from demo_models import face68_model, body_seg_model, pp_matting_model


def main(origin_image: Image, template_image: Image, maas_url, maas_token) -> Image:
    # 设置数据缓存目录
    _save_dir = generate_random_dir()
    # 数据先保存
    origin_filename = generate_random_file(dir_=_save_dir)
    origin_image.save(origin_filename)
    template_filename = generate_random_file(dir_=_save_dir)
    template_image.save(template_filename)
    # 调模型
    try:
        model = PortraitSolution(
            body_seg_model, face68_model, pp_matting_model)
        filename = model(image_path=origin_filename,
                         template_path=template_filename, save_dir=_save_dir)
        final_mask_filename = os.path.join(_save_dir, "neck-mask.png")
        # image = Image.open(filename)
        # maas
        maas = MaaS(endpoint=maas_url, token=maas_token)
        init_image = maas.upload_file(filename)
        mask_image = maas.upload_file(final_mask_filename)  # 上传合并图
        image = Image.open(BytesIO(requests.get(maas.infer_with_result(
            init_image=init_image, mask_image=mask_image, image_num=1)[0]).content))
    except Exception as e:
        logger.exception(str(e))
        image = None
    os.system(f"rm -rf {_save_dir}")
    return image


def demo():
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                origin_image = gr.Image(type="pil", height=480, label="原始图像")
                template_image = gr.Image(
                    type="pil", height=480, label="证件参考图")
            maas_url = gr.Textbox(
                value="https://lumicanvas.capitalonline.net", label="MaaS 平台 地址")
            maas_token = gr.Textbox(label="MaaS 平台 Token")
        out_image = gr.Image(
            type="pil")
    gr.Button("运行").click(main, inputs=[
        origin_image, template_image, maas_url, maas_token], outputs=[out_image])
